import cv2
import time

from yolo import YOLO
from PIL import Image
import numpy as np

video_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv']  # 可检测视频格式

data_path = "data/test.mp4"

# -------------------------------------------------------------------------#
#   crop                指定了是否在单张图片预测后对目标进行截取
#   count               指定了是否进行目标的计数
# -------------------------------------------------------------------------#
crop = False
count = False

cuda = True

data_format = data_path.split(".")[-1]

# 检测视频
video_save_path = ""
if data_format in video_formats:
    model = YOLO(cuda=cuda)
    capture = cv2.VideoCapture(data_path)

    _, frame = capture.read()

    fps = 0.0
    # noinspection PyRedundantParentheses
    while True:
        t1 = time.time()
        # 读取某一帧
        _, frame = capture.read()

        # 格式转变，BGRtoRGB
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 转变成Image
        frame = Image.fromarray(np.uint8(frame))
        # 进行检测
        frame = np.array(model.detect_image(frame))
        # RGBtoBGR满足opencv显示格式
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

        fps = (fps + (1. / (time.time() - t1))) / 2
        # noinspection PyRedundantParentheses
        # print("fps= %.2f" % (fps))
        # noinspection PyRedundantParentheses
        frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow("video", frame)
        c = cv2.waitKey(1) & 0xff

        if c == 27:
            capture.release()
            break

    print("视频检测完毕！！！")
    capture.release()
    cv2.destroyAllWindows()
else:
    print("视频格式或者视频本身有问题，无法检测！！！")
